package utils

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/url"
	"tuiwen-go/internal/config"

	"github.com/gorilla/websocket"
)

// PromptRequest 结构体用于封装发送到 ComfyUI 的请求
type PromptRequest struct {
	Prompt   map[string]interface{} `json:"prompt"`
	ClientID string                 `json:"client_id"`
}

// QueuePromptResponse 结构体用于解析 queue_prompt 的响应
type QueuePromptResponse struct {
	PromptID string                 `json:"prompt_id"`
	Error    map[string]interface{} `json:"error"`
}

// HistoryResponse 结构体用于解析 get_history 的响应
type HistoryResponse map[string]struct {
	Outputs map[string]struct {
		Images []map[string]interface{} `json:"images,omitempty"`
		Audio  []map[string]interface{} `json:"audio,omitempty"`
		Data   map[string]interface{}   `json:"data,omitempty"`
	} `json:"outputs"`
}

// WebSocketMessage 结构体用于解析 WebSocket 消息
type WebSocketMessage struct {
	Type string `json:"type"`
	Data struct {
		Node     interface{} `json:"node"`
		PromptID string      `json:"prompt_id"`
	} `json:"data"`
}

// QueuePrompt 向服务器发送提示
func QueuePrompt(prompt map[string]interface{}) (*QueuePromptResponse, error) {
	serverAddress := config.LoadedConfig.ComfyUI.ServerAddress
	clientID := config.LoadedConfig.ComfyUI.ClientID
	p := PromptRequest{
		Prompt:   prompt,
		ClientID: clientID,
	}
	data, err := json.Marshal(p)
	if err != nil {
		return nil, err
	}

	req, err := http.NewRequest("POST", fmt.Sprintf("http://%s/prompt", serverAddress), bytes.NewBuffer(data))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/json")

	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	var queueResp QueuePromptResponse
	if err := json.Unmarshal(body, &queueResp); err != nil {
		return nil, err
	} else {
		if queueResp.Error != nil {
			errorData, err := json.Marshal(queueResp.Error)
			if err != nil {
				return nil, fmt.Errorf("failed to marshal error data: %v", err)
			}
			return nil, errors.New(string(errorData))
		}
	}

	return &queueResp, nil
}

// GetImage 从服务器下载图像数据
func GetImage(filename, subfolder, folderType string) ([]byte, error) {
	serverAddress := config.LoadedConfig.ComfyUI.ServerAddress
	data := url.Values{
		"filename":  {filename},
		"subfolder": {subfolder},
		"type":      {folderType},
	}
	urlString := fmt.Sprintf("http://%s/view?%s", serverAddress, data.Encode())

	resp, err := http.Get(urlString)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	return io.ReadAll(resp.Body)
}

// GetHistory 获取历史记录
func GetHistory(promptID string) (HistoryResponse, error) {
	serverAddress := config.LoadedConfig.ComfyUI.ServerAddress
	urlString := fmt.Sprintf("http://%s/history/%s", serverAddress, promptID)
	resp, err := http.Get(urlString)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	var history HistoryResponse
	if err := json.Unmarshal(body, &history); err != nil {
		return nil, err
	}

	return history, nil
}

// GetOutputData 通过 WebSocket 接收消息并下载数据
func GetOutputData(prompt map[string]interface{}, outputType string) (map[string][][]byte, string, error) {
	serverAddress := config.LoadedConfig.ComfyUI.ServerAddress
	clientID := config.LoadedConfig.ComfyUI.ClientID
	u := "ws://" + serverAddress + "/ws?clientId=" + clientID
	conn, _, err := websocket.DefaultDialer.Dial(u, nil)
	if err != nil {
		return nil, "", err
	}
	defer conn.Close()

	queueResp, err := QueuePrompt(prompt)
	if err != nil {
		return nil, "", err
	}
	promptID := queueResp.PromptID

	outputData := make(map[string][][]byte)

	history, err := GetHistory(promptID)
	log.Printf("history: %v", history)
	if err != nil || history == nil || len(history) == 0 {
		log.Printf("开始监听 WebSocket 消息，等待执行完成信号，PromptID: %s", promptID)
		for {
			_, messageBytes, err := conn.ReadMessage()
			if err != nil {
				log.Printf("读取 WebSocket 消息出错: %v", err)
				return nil, "", err
			}

			var message WebSocketMessage
			if err := json.Unmarshal(messageBytes, &message); err != nil {
				log.Printf("解析 WebSocket 消息出错: %v，跳过该消息", err)
				continue
			}
			log.Printf("收到 WebSocket 消息，类型: %s，Node: %v，PromptID: %s", message.Type, message.Data.Node, message.Data.PromptID)

			if message.Type == "executing" {
				if message.Data.Node == nil && message.Data.PromptID == promptID {
					log.Printf("检测到执行完成信号，PromptID: %s，退出监听", promptID)
					break
				}
			}
		}
		history, err = GetHistory(promptID)
		if err != nil {
			return nil, "", err
		}
	}
	var filename = ""
	if hist, ok := history[promptID]; ok {
		for nodeID, nodeOutput := range hist.Outputs {
			var outputFiles []map[string]interface{}
			switch outputType {
			case "images":
				outputFiles = nodeOutput.Images
			case "audio":
				outputFiles = nodeOutput.Audio
			default:
				continue
			}

			if len(outputFiles) > 0 {
				var filesOutput [][]byte
				for _, file := range outputFiles {
					filename, _ = file["filename"].(string)
					subfolder, _ := file["subfolder"].(string)
					fileType, _ := file["type"].(string)
					fileData, err := GetImage(filename, subfolder, fileType)
					if err != nil {
						log.Printf("Error getting file %s: %v", filename, err)
						continue
					}
					filesOutput = append(filesOutput, fileData)
				}
				outputData[nodeID] = filesOutput
			}
		}
	}

	return outputData, filename, nil
}

// GetImages 通过 WebSocket 接收消息并下载图像
func GetImages(prompt map[string]interface{}) (map[string][][]byte, string, error) {
	return GetOutputData(prompt, "images")
}

// GetAudios 通过 WebSocket 接收消息并下载音频
func GetAudios(prompt map[string]interface{}) (map[string][][]byte, string, error) {
	return GetOutputData(prompt, "audio")
}
